#include "bwt.h"
#include <stack>
#include <iostream>
#include <cstdlib>
#include <vector>

using namespace std;
using namespace bwt;

BWTencoder::BWTencoder(int32_t buf_size): _buf_size(buf_size) {
    _A = new int32_t[_buf_size];
}

BWTencoder::~BWTencoder() {
    delete[] _A;
}

void BWTencoder::encode(hzstream* hz)
{
    hz->out_size = hz->in_size + sizeof(int32_t);
    hz->out_buf = new uint8_t[hz->out_size];
    (reinterpret_cast<int32_t*>(hz->out_buf))[0]=
    this->encode(hz->in_buf,
                (hz->out_buf + sizeof(int32_t)),
                 hz->in_size);
}

int32_t BWTencoder::encode(uint8_t* in_buf, uint8_t* out_buf,int32_t buf_size)
{
    if(buf_size==1)
    {
        out_buf[0]=in_buf[0];
        return -1;
    }

    int32_t idx=0;

    _sortSuffixes(in_buf,_A,buf_size);

    for(int32_t i=0;i<buf_size;i++)
    {
        if(_A[i]==0)
            out_buf[i]=in_buf[buf_size-1];
        else
            out_buf[i]=in_buf[_A[i]-1];
        if(_A[i]==1)
            idx = i;
    }
    
    return idx;
}

void BWTencoder::_radixSort(uint8_t* in_buf,int32_t buf_size,int32_t* A,int32_t start,int32_t end)
{
    int32_t* tmp = new int32_t[end-start+1];
    stack<classForRadix> stackForRadix;
    stackForRadix.push(classForRadix(1,start,end));

    int32_t dist;
    int32_t s,e;
    int32_t bucket[257];
    int32_t index[257];
    classForRadix c;

    while(!stackForRadix.empty())
    {
        c=stackForRadix.top();
        stackForRadix.pop();

        dist=c.dist;
        s=c.start;
        e=c.end;

        for(int32_t i=0;i<257;i++)
        {
            bucket[i]=0;
        }

        for(int32_t i=s;i<e;i++)//count elements
        {
            if(dist==buf_size)//consider $(end of string, just to protect from looping)
            {
                bucket[256]++;
            }
            else
            {
                bucket[in_buf[(A[i]+dist)%buf_size]]++;
            }
        }

        index[256]=0;
        index[0]=bucket[256];
        for(int32_t i=1;i<256;i++)
        {
          index[i]=index[i-1]+bucket[i-1];
        }

        for(int32_t i=s;i<e;i++)
        {
            if(dist==buf_size)
            {
                tmp[index[256]]=A[i];
                index[256]++;
            }
            else
            {
                tmp[index[in_buf[(A[i]+dist)%buf_size]]]=A[i];
                index[in_buf[(A[i]+dist)%buf_size]]++;
            }
        }

        for(int32_t i=s,j=0;i<e;i++,j++)
            A[i]=tmp[j];

        for(int32_t i=256;i>=0;i--)//this direction, because it's a stack
        {
            if(bucket[i]>1)
            {
                //cout<<"I must sort "<<static_cast<char>(i)<<' '<<index[i]-bucket[i]+s<<' '<<index[i]+s<<endl;
                if(dist+1<buf_size)
                    stackForRadix.push(classForRadix(dist+1,index[i]-bucket[i]+s,index[i]+s));
            }
        }
    }

    delete[] tmp;

}

void BWTencoder::_sortSuffixes(uint8_t* in_buf, int32_t* A,int32_t buf_size)
{
    int32_t bucketSizeX[256],bucketSizeY[256],bucketPtrX[256],bucketPtrY[256];

    for(int32_t i=0;i<buf_size;i++)
        A[i]=-1;

    for(int32_t i=0;i<256;i++)//step 1-1
    {
        bucketSizeX[i]=0;
        bucketSizeY[i]=0;
    }

    //buf_size--;
    for(int32_t i=0;i<buf_size;i++)//step 1-2
    {///e:eae
        if(in_buf[i]>in_buf[(i+1)%buf_size])
        {
            bucketSizeX[in_buf[i]]++;
        }
        else
        {
            bucketSizeY[in_buf[i]]++;
        }
    }
    /*buf_size++;
    bucketSizeY[in_buf[buf_size-1]]++;*/

    //step 1-3:
    bucketPtrX[0]=0;
    bucketPtrY[0]=0+bucketSizeX[0];
    for(int32_t i=1;i<256;i++)
    {
        bucketPtrX[i]=bucketPtrY[i-1]+bucketSizeY[i-1];
        bucketPtrY[i]=bucketPtrX[i]+bucketSizeX[i];
    }

    //step1-4
    //buf_size--;
    for(int32_t i=0,j=0;i<buf_size;i++)
    {
        j=i+1;
        if(i+1==buf_size)
            j=0;

        if(in_buf[i]<=in_buf[j])//in Y
        {
            A[bucketPtrY[in_buf[i]]]=i;
            bucketPtrY[in_buf[i]]++;
        }
    }
    /*buf_size++;

    A[bucketPtrY[in_buf[buf_size-1]]]=buf_size-1;
    bucketPtrY[in_buf[buf_size-1]]++;*/

    for(int32_t i=0;i<256;i++)//step 2
    {
        if(bucketSizeY[i]>1)
        {
            //sort from bucketPtrY[i]-bucketSizeY[i]+1 to bucketPtrY[i]
            //cout<<"sort from "<<bucketPtrY[i]-bucketSizeY[i]+1<<" to "<<bucketPtrY[i]<<endl;
            _radixSort(in_buf,buf_size,A,(bucketPtrY[i]-bucketSizeY[i]),bucketPtrY[i]);
        }
    }

    for(int32_t i=0;i<buf_size;i++)//step 3
    {
        if(A[i]>0)
        {
            //cout<<A[i]<<' '<<in_buf[A[i]]<<endl;
            if(in_buf[A[i]-1]>in_buf[A[i]])//A[i]-1 in X
            {
                A[bucketPtrX[in_buf[A[i]-1]]]=A[i]-1;
                bucketPtrX[in_buf[A[i]-1]]++;
            }
        }
        else//A[i]==0
        {
            if(in_buf[buf_size-1]>in_buf[A[i]])//A[i]-1 in X
            {
                A[bucketPtrX[in_buf[buf_size-1]]]=buf_size-1;
                bucketPtrX[in_buf[buf_size-1]]++;
            }
        }
    }

}

BWTdecoder::BWTdecoder(int32_t buf_size): _buf_size(buf_size) {
    _VT = new int32_t[_buf_size];
}

BWTdecoder::~BWTdecoder() {
    delete[] _VT;
}

void BWTdecoder::decode(hzstream* hz)
{
    hz->out_size = hz->in_size - sizeof(int32_t);
    hz->out_buf = new uint8_t[hz->out_size];
    this->decode((hz->in_buf + sizeof(int32_t)),
                 (reinterpret_cast<int32_t*>(hz->in_buf))[0],
                 hz->out_buf,
                 hz->out_size);
}

void BWTdecoder::decode(uint8_t* in_buf,int32_t index,uint8_t* out_buf, int32_t buf_size)
{
    if(buf_size==1)
    {
        out_buf[0]=in_buf[0];
        return;
    }

    int32_t lookup[256];
    int32_t idx=index;

    _makeLookupInVP(in_buf,lookup,buf_size);

    //cout<<"A:"<<endl;
    for(int32_t i=0;i<buf_size;i++)
    {
        _VT[ lookup[ in_buf[i] ]++ ]=i;
    }

    for(int32_t i=0;i<buf_size;i++)
    {
        out_buf[i]=in_buf[idx];
        idx=_VT[idx];
    }

}

void BWTdecoder::_makeLookupInVP(uint8_t *in_buf,int32_t* lookup,int32_t buf_size)
{
    int32_t bucket[256];

    for(int32_t i=0;i<256;i++)
    {
        bucket[i]=0;
    }

    for(int32_t i=0;i<buf_size;i++)//count elements
    {
        bucket[in_buf[i]]++;
    }

    lookup[0]=0;
    for(int32_t i=1;i<256;i++)
    {
        lookup[i]=lookup[i-1]+bucket[i-1];
    }
}
